#ifndef __EXTERN_H
#define __EXTERN_H
#include "chatglm.h"
// #include "chatglm_cpp_sys/src/lib.rs.h"
#include "rust/cxx.h"

namespace chatglm
{
    
    class ExternStreamer : public BaseStreamer
    {
    public:
        ExternStreamer();
        ExternStreamer(BaseTokenizer *tokenizer)
            : tokenizer_(tokenizer), is_prompt_(true), is_first_line_(true), print_len_(0)
        {
        }
        void put(const std::vector<int> &output_ids) override;
        void set_callback(rust::Fn<void(rust::Str s)> process_callback, rust::Fn<void()> end_callback);
        void end() override;

    private:
        rust::Fn<void(rust::Str s)> _process_callback;
        rust::Fn<void()> _end_callback;
        BaseTokenizer *tokenizer_;
        bool is_prompt_;
        bool is_first_line_;
        std::vector<int> token_cache_;
        int print_len_;
    };
    class ExternPipeline
    {
    private:
        ExternStreamer *_extern_streamer;
        GenerationConfig _gen_config;
        Pipeline *_pipeline;

    public:
        ExternPipeline(const std::string &path, GenerationConfig gen_config)
        {
            _pipeline = new Pipeline(path);
            _extern_streamer = new ExternStreamer(_pipeline->tokenizer.get());
            _gen_config = gen_config;
        }
        ~ExternPipeline()
        {
            delete _pipeline;
            delete _extern_streamer;
        }
        void set_callback(rust::Fn<void(rust::Str s)> process_callback, rust::Fn<void()> end_callback) const;
        rust::String generate_string(const rust::Str prompt,
                                     bool use_callback = false) const;
    };
    std::unique_ptr<GenerationConfig> generation_config_new(int max_length = 2048, int max_new_tokens = -1, int max_context_length = 512,
                                                            bool do_sample = true, int top_k = 0, float top_p = 0.7, float temperature = 0.95,
                                                            float repetition_penalty = 1.f, int num_threads = 0);
    std::unique_ptr<ExternPipeline> extern_pipeline_new(rust::Str path, std::unique_ptr<GenerationConfig> gen_config);
}
#endif